import numpy as np
import faiss

class FaissKNeighbors:
    def __init__(self, k=1, res=None):
        self.index = None
        self.y = None
        self.k = k
        self.res = res

    def fit(self, X, y):
        self.index = faiss.IndexFlatL2(X.shape[1])  # 初始化为IndexFlatL2索引
        if self.res is not None:
            self.index = faiss.index_cpu_to_gpu(self.res, 0, self.index)  # 将索引转移到GPU上
        self.index.add(X.astype(np.float32))  # 将训练数据加入索引
        self.y = y

    def predict(self, X):
        distances, indices = self.index.search(X.astype(np.float32), self.k)
        votes = self.y[indices]
        predictions = np.array([np.argmax(np.bincount(vote)) for vote in votes])
        return predictions

    def score(self, X, y):
        predictions = self.predict(X)
        accuracy = np.mean(predictions == y)
        return accuracy
